import numpy as np
import lightgbm as lgb
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

from sklearn.datasets import make_classification
import pandas as pd
import copy
demand_train_A = 'data/demand_train_A.csv'
demand_train_A = pd.read_csv(demand_train_A)
tmpgroup = None
p=0
for i, group in demand_train_A.groupby('unit'):
    p+=1
    if p >= 2:
        tmpgroup = copy.deepcopy(group)
        break
tmpts = pd.to_datetime(tmpgroup['ts'])
tmpgroup['month'] = tmpts.dt.month
tmpgroup['quarter'] = tmpts.dt.quarter
tmpgroup['dayofweek'] = tmpts.dt.dayofweek

# 构造特征的是月份、季节、星期三个特征的独热编码
month_list = np.eye(12)
quarter_list = np.eye(4)
dayofweek_list = np.eye(7)
month_dict = {}
quarter_dict = {}
dayofweek_dict = {}
for i in range(len(month_list)):
    month_dict[i+1] = month_list[i]
for i in range(len(quarter_list)):
    quarter_dict[i+1] = quarter_list[i]
for i in range(len(dayofweek_list)):
    dayofweek_dict[i] = dayofweek_list[i]

pdx_map = {
    'month_list':[],
    'quarter_list':[],
    'dayofweek_list':[],
    'one_hots':[]
}
pdx_map['month_list'] = tmpgroup['month'].to_list()
pdx_map['quarter_list'] = tmpgroup['quarter'].to_list()
pdx_map['dayofweek_list'] = tmpgroup['dayofweek'].to_list()
for i in range(len(pdx_map['month_list'])):
    tmp = np.concatenate((month_dict[pdx_map['month_list'][i]], quarter_dict[pdx_map['quarter_list'][i]], dayofweek_dict[pdx_map['dayofweek_list'][i]]))
    pdx_map['one_hots'].append(tmp)
pdx_map['one_hots'] = np.array(pdx_map['one_hots'])
    
    

x = pdx_map['one_hots']
y = tmpgroup['qty'].to_list()
x_train, x_test, y_train, y_test = train_test_split(x, y)
lgb_train = lgb.Dataset(x_train, y_train)
lgb_eval = lgb.Dataset(x_test, y_test, reference=lgb_train)

# 将参数写成字典下形式
params = {
    'task': 'train',
    'boosting_type': 'gbdt',  # 设置提升类型
    'objective': 'regression', # 目标函数
    'metric': {'l2', 'auc'},  # 评估函数
    'num_leaves': 31,   # 叶子节点数
    'learning_rate': 0.05,  # 学习速率
    'feature_fraction': 0.9, # 建树的特征选择比例
    'bagging_fraction': 0.8, # 建树的样本采样比例
    'bagging_freq': 5,  # k 意味着每 k 次迭代执行bagging
    'verbose': 1 # <0 显示致命的, =0 显示错误 (警告), >0 显示信息
}
 
print('Start training...')



gbm = lgb.train(
    params,
    lgb_train,
    num_boost_round=20,
    valid_sets=lgb_eval,
    early_stopping_rounds=5) # 训练数据需要参数列表和数据集
 
print('Save model...') 
 
gbm.save_model('model.txt')   # 训练后保存模型到文件

# 预测数据集
y_pred = gbm.predict(x_test, num_iteration=gbm.best_iteration) #如果在训练期间启用了早期停止，可以通过best_iteration方式从最佳迭代中获得预测
# 评估模型
print('The rmse of prediction is:', mean_squared_error(y_test, y_pred) ** 0.5) # 计算真实值和预测值之间的均方根误差


